//
// Created by Ivor_ on 2023/5/8.
//

#include "diagonalization.h"

double* dJacobiD(int size, double** matrix, double** eigVecs, double eps) {
    double* eigVals = (double*)calloc(size, sizeof(double));
    double** iterMat = (double**)calloc(size, sizeof(double*));
    for (int i = 0; i < size; ++i) {
        iterMat[i] = (double*)calloc(size, sizeof(double));
        for (int j = 0; j < size; ++j) {
            iterMat[i][j] = matrix[i][j];
            eigVecs[i][j] = 1.*(i == j);
        }
    }
    double** tempSpace = (double**)calloc(4, sizeof(double*));
    for (int i = 0; i < 4; ++i) {
        tempSpace[i] = (double*)calloc(size, sizeof(double));
    }
    while (1) {
        int loc[2] = {0, 1};
        double maxValue = iterMat[loc[0]][loc[1]];
        for (int i = 0; i < size-1; ++i) {
            for (int j = i+1; j < size; ++j) {
                if (fabs(iterMat[i][j]) > fabs(maxValue)) {
                    maxValue = iterMat[i][j];
                    loc[0] = i;
                    loc[1] = j;
                }
            }
        }
        if (fabs(maxValue) < eps) {
            break;
        }
        double tempB = (iterMat[loc[1]][loc[1]] - iterMat[loc[0]][loc[0]]) / iterMat[loc[0]][loc[1]];
        double tanA1 = (- tempB + sqrt(tempB * tempB + 4.)) / 2.;
        double tanA2 = (- tempB - sqrt(tempB * tempB + 4.)) / 2.;
        double tanA = fabs(tanA1) > fabs(tanA2) ? tanA2 : tanA1;
        double cosA = 1. / sqrt(tanA * tanA + 1.);
        double sinA = tanA / sqrt(tanA * tanA + 1.);
        for (int i = 0; i < size; ++i) {
            tempSpace[0][i] = eigVecs[i][loc[0]];
            tempSpace[1][i] = eigVecs[i][loc[1]];
        }
        for (int i = 0; i < size; ++i) {
            eigVecs[i][loc[0]] = tempSpace[0][i]*cosA - tempSpace[1][i]*sinA;
            eigVecs[i][loc[1]] = tempSpace[1][i]*cosA + tempSpace[0][i]*sinA;
        }
        for (int i = 0; i < size; ++i) {
            tempSpace[0][i] = iterMat[loc[0]][i];
            tempSpace[1][i] = iterMat[loc[1]][i];
            tempSpace[2][i] = iterMat[i][loc[0]];
            tempSpace[3][i] = iterMat[i][loc[1]];
        }
        for (int i = 0; i < size; ++i) {
            iterMat[loc[0]][i] = tempSpace[0][i]*cosA - tempSpace[1][i]*sinA;
            iterMat[loc[1]][i] = tempSpace[1][i]*cosA + tempSpace[0][i]*sinA;
            iterMat[i][loc[0]] = tempSpace[2][i]*cosA - tempSpace[3][i]*sinA;
            iterMat[i][loc[1]] = tempSpace[3][i]*cosA + tempSpace[2][i]*sinA;
        }
        iterMat[loc[0]][loc[0]] = cosA*(tempSpace[0][loc[0]]*cosA - tempSpace[1][loc[0]]*sinA);
        iterMat[loc[0]][loc[0]] -= sinA*(tempSpace[0][loc[1]]*cosA - tempSpace[1][loc[1]]*sinA);
        iterMat[loc[0]][loc[1]] = sinA*(tempSpace[0][loc[0]]*cosA - tempSpace[1][loc[0]]*sinA);
        iterMat[loc[0]][loc[1]] += cosA*(tempSpace[0][loc[1]]*cosA - tempSpace[1][loc[1]]*sinA);
        iterMat[loc[1]][loc[0]] = cosA*(tempSpace[0][loc[0]]*sinA + tempSpace[1][loc[0]]*cosA);
        iterMat[loc[1]][loc[0]] -= sinA*(tempSpace[1][loc[1]]*cosA + tempSpace[0][loc[1]]*sinA);
        iterMat[loc[1]][loc[1]] = sinA*(tempSpace[0][loc[0]]*sinA + tempSpace[1][loc[0]]*cosA);
        iterMat[loc[1]][loc[1]] += cosA*(tempSpace[1][loc[1]]*cosA + tempSpace[0][loc[1]]*sinA);
    }
    for (int i = 0; i < size; ++i) {
        eigVals[i] = iterMat[i][i];
        free(iterMat[i]);
    }
    for (int i = 0; i < 4; ++i) {
        free(tempSpace[i]);
    }
    free(iterMat);
    free(tempSpace);
    return eigVals;
}

_Complex double* cJacobiD(int size, _Complex double** matrix, _Complex double** eigVecs, double eps) {
    _Complex double* eigVals = (_Complex double*)calloc(size, sizeof(_Complex double));
    _Complex double** iterMat = (_Complex double**)calloc(size, sizeof(_Complex double*));
    for (int i = 0; i < size; ++i) {
        iterMat[i] = (_Complex double*)calloc(size, sizeof(_Complex double));
        for (int j = 0; j < size; ++j) {
            iterMat[i][j] = matrix[i][j];
            eigVecs[i][j] = 1.*(i == j);
        }
    }
    _Complex double** tempSpace = (_Complex double**)calloc(4, sizeof(_Complex double*));
    for (int i = 0; i < 4; ++i) {
        tempSpace[i] = (_Complex double*)calloc(size, sizeof(_Complex double));
    }
    while (1) {
        int loc[2] = {0, 1};
        _Complex double maxValue = iterMat[loc[0]][loc[1]];
        for (int i = 0; i < size-1; ++i) {
            for (int j = i+1; j < size; ++j) {
                if (cabs(iterMat[i][j]) > cabs(maxValue)) {
                    maxValue = iterMat[i][j];
                    loc[0] = i;
                    loc[1] = j;
                }
            }
        }
        if (cabs(maxValue) < eps) {
            break;
        }
        double tempB = (creal(iterMat[loc[1]][loc[1]]) - creal(iterMat[loc[0]][loc[0]])) / cabs(iterMat[loc[0]][loc[1]]);
        double tanA1 = (- tempB + sqrt(tempB * tempB + 4.)) / 2.;
        double tanA2 = (- tempB - sqrt(tempB * tempB + 4.)) / 2.;
        double tanA = fabs(tanA1) > fabs(tanA2) ? tanA2 : tanA1;
        double cosA = 1. / sqrt(tanA * tanA + 1.);
        double sinA = tanA / sqrt(tanA * tanA + 1.);
        double theta = .5 * carg(iterMat[loc[0]][loc[1]]);
        for (int i = 0; i < size; ++i) {
            tempSpace[0][i] = eigVecs[i][loc[0]];
            tempSpace[1][i] = eigVecs[i][loc[1]];
        }
        for (int i = 0; i < size; ++i) {
            eigVecs[i][loc[0]] = tempSpace[0][i]*cosA*cexp(_Complex_I*theta) - tempSpace[1][i]*sinA*cexp(-_Complex_I*theta);
            eigVecs[i][loc[1]] = tempSpace[1][i]*cosA*cexp(-_Complex_I*theta) + tempSpace[0][i]*sinA*cexp(_Complex_I*theta);
        }
        for (int i = 0; i < size; ++i) {
            tempSpace[0][i] = iterMat[loc[0]][i];
            tempSpace[1][i] = iterMat[loc[1]][i];
            tempSpace[2][i] = iterMat[i][loc[0]];
            tempSpace[3][i] = iterMat[i][loc[1]];
        }
        for (int i = 0; i < size; ++i) {
            iterMat[loc[0]][i] = tempSpace[0][i]*cosA*cexp(-_Complex_I*theta) - tempSpace[1][i]*sinA*cexp(_Complex_I*theta);
            iterMat[loc[1]][i] = tempSpace[1][i]*cosA*cexp(_Complex_I*theta) + tempSpace[0][i]*sinA*cexp(-_Complex_I*theta);
            iterMat[i][loc[0]] = tempSpace[2][i]*cosA*cexp(_Complex_I*theta) - tempSpace[3][i]*sinA*cexp(-_Complex_I*theta);
            iterMat[i][loc[1]] = tempSpace[3][i]*cosA*cexp(-_Complex_I*theta) + tempSpace[2][i]*sinA*cexp(_Complex_I*theta);
        }
        iterMat[loc[0]][loc[0]] = cosA*cexp(_Complex_I*theta)*(tempSpace[0][loc[0]]*cosA*cexp(-_Complex_I*theta) - tempSpace[1][loc[0]]*sinA*cexp(_Complex_I*theta));
        iterMat[loc[0]][loc[0]] -= sinA*cexp(-_Complex_I*theta)*(tempSpace[0][loc[1]]*cosA*cexp(-_Complex_I*theta) - tempSpace[1][loc[1]]*sinA*cexp(_Complex_I*theta));
        iterMat[loc[0]][loc[1]] = sinA*cexp(_Complex_I*theta)*(tempSpace[0][loc[0]]*cosA*cexp(-_Complex_I*theta) - tempSpace[1][loc[0]]*sinA*cexp(_Complex_I*theta));
        iterMat[loc[0]][loc[1]] += cosA*cexp(-_Complex_I*theta)*(tempSpace[0][loc[1]]*cosA*cexp(-_Complex_I*theta) - tempSpace[1][loc[1]]*sinA*cexp(_Complex_I*theta));
        iterMat[loc[1]][loc[0]] = cosA*cexp(_Complex_I*theta)*(tempSpace[0][loc[0]]*sinA*cexp(-_Complex_I*theta) + tempSpace[1][loc[0]]*cosA*cexp(_Complex_I*theta));
        iterMat[loc[1]][loc[0]] -= sinA*cexp(-_Complex_I*theta)*(tempSpace[1][loc[1]]*cosA*cexp(_Complex_I*theta) + tempSpace[0][loc[1]]*sinA*cexp(-_Complex_I*theta));
        iterMat[loc[1]][loc[1]] = sinA*cexp(_Complex_I*theta)*(tempSpace[0][loc[0]]*sinA*cexp(-_Complex_I*theta) + tempSpace[1][loc[0]]*cosA*cexp(_Complex_I*theta));
        iterMat[loc[1]][loc[1]] += cosA*cexp(-_Complex_I*theta)*(tempSpace[1][loc[1]]*cosA*cexp(_Complex_I*theta) + tempSpace[0][loc[1]]*sinA*cexp(-_Complex_I*theta));
    }
    for (int i = 0; i < size; ++i) {
        eigVals[i] = iterMat[i][i];
        free(iterMat[i]);
    }
    for (int i = 0; i < 4; ++i) {
        free(tempSpace[i]);
    }
    free(iterMat);
    free(tempSpace);
    return eigVals;
}
